import matplotlib.pyplot as plt
import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

batch_size = 256
lr = 5e-3
epochs = 5

# 设置数据变换
from torchvision import transforms

image_size = 28
data_transform = transforms.Compose([
    # transforms.ToPILImage(),   # 这一步取决于后续的数据读取方式，如果使用内置数据集则不需要
    transforms.Resize(image_size),
    transforms.ToTensor()
])

## 使用torchvision自带数据集，下载可能需要一段时间
from torchvision import datasets

# The dataset will be downloaded to ROOT_DIR
if os.path.isdir('D:/'):
    ROOT_DIR = 'D:/'
else:
    ROOT_DIR = '.'
    
train_data = datasets.FashionMNIST(root=ROOT_DIR, train=True, download=True, transform=data_transform)
test_data = datasets.FashionMNIST(root=ROOT_DIR, train=False, download=True, transform=data_transform)
classes = np.array(["T-shirt/top","Trouser","Pullover","Dress","Coat","Sandal","Shirt","Sneaker","Bag","Ankle boot"])

train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=True)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

# %matplotlib inline

image, label = next(iter(train_loader))
print(image.shape, label.shape)

for i in range(5):
    plt.figure(figsize=(2,2))
    plt.imshow(image[i][0], cmap="gray")
    plt.title(label[i].item())
    plt.axis('off')
    plt.show()

class MyCNN(nn.Module):
    def __init__(self):
        super(MyCNN, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(16, 32, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(32 * 4 * 4, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        return self.layers(x)
    # YOUR CODE HERE
    # raise NotImplementedError()
    
model = MyCNN().to(device) 
model

assert len(model.layers) == 10 # 共10层
assert isinstance( model.layers[0], nn.Conv2d) # 第一层为卷积层
assert model.layers[0].kernel_size == (5,5) # 第一层卷积尺为5x5
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

model.train()

loss_history = []
acc_history = []

# Run the training
for epoch in range(epochs):
    running_loss = 0.0
    correct = 0
    total = 0

    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / len(train_loader)
    epoch_acc = correct / total

    loss_history.append(epoch_loss)
    acc_history.append(epoch_acc)

    print(f"Epoch {epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.4f}")

    # YOUR CODE HERE
    # raise NotImplementedError()

plt.plot(loss_history)
plt.plot(acc_history)

import matplotlib.pyplot as plt

# 从dataloader中选取max_samples个样本进行分类评价，返回准确率
def evaluate_rgb_images(model, dataloader, max_samples = 100):

    model.eval()    
    sample_cnt = 0
    correct = 0
    for inputs, labels in dataloader:
    
        if sample_cnt >= max_samples:
            break
        sample_cnt += len(labels)
        
        plt.figure(figsize=(10,len(inputs)/10+1))
        for idx, img  in enumerate(inputs):
            plt.subplot(int(len(inputs)/10)+1, 10, idx+1) 
            plt.imshow(img[0])
            plt.axis('off')
        plt.show()
        # Move the inputs and labels to the device
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        print('truth: \n', labels)
        print('prediction: \n', preds)
        correct += (preds == labels).float().sum()
        
    accuracy = 100 * correct / sample_cnt
    print('classification accuracy (%) =', accuracy.item())
    return accuracy

test_acc = evaluate_rgb_images(model, test_loader, max_samples = 200)
assert test_acc > 0.85 # 测试准确率应在85%以上
Student_id = '2335040120' # ！！！这里填你的学号

save_path = Student_id +".pt" # 模型保存为 学号.pt
torch.save(model, save_path)
assert os.path.isfile(save_path)
model = torch.load(save_path)
assert len(model.layers) == 10